# -*- coding: UTF-8 -*-

import pymongo
import traceback

"""
    功能：etl脚本
         读取路段静态数据文件，并匹配路段与rtic的对应关系
"""

db_name = "baotou_traffic"
roadinfo_col = "roadinfo"

def read_roadinfo():

    # with open("roadinfo927.txt", "rb") as reader, pymongo.MongoClient("192.168.4.128", 27017) as conn:
    with open("roadinfo927.txt", "rb") as reader, pymongo.MongoClient("123.207.109.239", 27017) as conn:

        col = conn[db_name][roadinfo_col]

        bulk = []
        total = 0
        for line in reader.readlines():

            line = line.replace("\r\n", "")
            data = line.split("|")

            roadinfo = {}
            roadinfo["road_id"] = data[0]
            roadinfo["road_name"] = data[1]
            roadinfo["road_kind"] = data[2]
            roadinfo["road_start"] = data[3]
            roadinfo["road_end"] = data[4]

            bulk.append(roadinfo)
            total += 1

            if total % 500 == 0:
                col.insert(bulk)
                bulk = []

        if len(bulk) > 0:
            col.insert(bulk)

        print "total:%s" % total

def roadrtic():
    with open("roadrtic.txt", "rb") as reader, pymongo.MongoClient("123.207.109.239", 27017) as conn:

        try:
            col = conn[db_name][roadinfo_col]

            road_rtic = {}
            for line in reader.readlines():
                try:
                    line = line.replace("\r\n", "")
                    data = line.split(",")

                    road_id = data[0]
                    rtic_id = data[1]

                    rtic_coordinate_seq = ""
                    for i in range(2, len(data), 1):
                        rtic_coordinate_seq += data[i]
                        rtic_coordinate_seq += ","
                    rtic_coordinate_seq = rtic_coordinate_seq[:len(rtic_coordinate_seq) - 1]

                    # print road_id, rtic_id, rtic_coordinate_seq
                    # print rtic_coordinate_seq

                    if road_id not in road_rtic:
                        road_rtic[road_id] = [{"rtic_id": rtic_id, "rtic_coordinate_seq": rtic_coordinate_seq}]
                    else:
                        road_rtic[road_id].append({"rtic_id": rtic_id, "rtic_coordinate_seq": rtic_coordinate_seq})

                except:
                    traceback.print_exc()

            # 存入mongo
            total = 0
            for road_id in road_rtic.keys():
                col.update({"road_id": road_id},
                           {"$set": {"road_rtic": road_rtic[road_id]}})
                total += 1

                if total % 500 == 0:
                    print "update count: %s" % total

            print "update total: %s" % total
        except:
            traceback.print_exc()

    print "process done!!"

def make_index():
    conn = pymongo.MongoClient("123.207.109.239", 27017)
    col = conn[db_name][roadinfo_col]
    col.create_index("road_id")
    col.create_index("road_kind")
    # col.create_index("road_rtic")
    # col.create_index("road_rtic.rtic_id")
    # col.create_index("road_rtic.rtic_coordinate_seq")

if __name__ == "__main__":
    read_roadinfo()
    roadrtic()
    make_index()



